import json
import os
from collections import deque
from typing import Optional, Tuple

import gym
import numpy as np
import pandas as pd
from gym import spaces
from gym.core import ObsType
from stable_baselines3 import A2C, DDPG, DQN, PPO, SAC, TD3

from backtest import BacktestModel, OrderType, PositionType
from utils import get_baostock_data
from view_backtrade_result import calculate_backtrade_result

MAX_ACCOUNT_BALANCE = 2147483647
MAX_NUM_SHARES = 2147483647
MAX_SHARE_PRICE = 5000
MAX_VOLUME = 1000e8
MAX_AMOUNT = 3e10
MAX_OPEN_POSITIONS = 5
MAX_STEPS = 20000
MAX_DAY_CHANGE = 1
OP_VOLUME = 1e5


class StockEnv(gym.Env):

    def __init__(self, stock_backtest: BacktestModel):
        super(StockEnv, self).__init__()
        self.stock_backtest = stock_backtest
        num_stocks = len(stock_backtest.datas)
        low = np.zeros([num_stocks], dtype=np.float32)
        high = np.zeros([num_stocks], dtype=np.float32) + 1
        self.action_space = spaces.Box(low=low, high=high, dtype=np.float32)

        self.observation_space = spaces.Box(
            low=-1, high=1, shape=self.observation().shape, dtype=np.float32
        )

        self.reset()

    def reset(
        self,
        *,
        seed: Optional[int] = None,
        options: Optional[dict] = None,
    ) -> Tuple[ObsType, dict]:
        self.done = False
        self.stock_backtest.reset()
        self.value_record = deque(maxlen=20)
        self.value_record.append(self.stock_backtest.get_value())
        self.all_value_record = []
        self.all_value_record.append(
            {
                "date": self.stock_backtest.get_cur_time(),
                "value": self.stock_backtest.get_value(),
            }
        )
        return self.observation(), {}

    def step(self, action):

        # each_stock_cash = action / action.sum() * before_value
        each_stock_name = list(self.stock_backtest.datas.keys())
        target_positions = {}
        sum_action = sum(action)
        now_value = self.stock_backtest.get_value()
        for i in range(len(action)):
            price = self.stock_backtest.get_data(each_stock_name[i], "close", 0)
            if not price:
                continue

            target_positions[each_stock_name[i]] = max(
                action[i] / sum_action * now_value / price - 10, 0
            )
        before_positions = self.stock_backtest.get_positions()

        modify_positions = {}

        for name in target_positions:
            if name in before_positions:
                before_size = before_positions[name].size
            else:
                before_size = 0
            modify_positions[name] = target_positions.get(name, 0) - before_size

        for name, size in modify_positions.items():
            if size < 0:
                try:
                    self.stock_backtest.order(OrderType(name, size))
                except Exception as e:
                    print(f"Deal order error: {e}")

        reward = None
        for name, size in modify_positions.items():
            if size > 0:
                try:
                    self.stock_backtest.order(OrderType(name, size))
                except Exception as e:
                    print(f"Deal order error: {e}")
                    # reward = -100

        self.done = self.stock_backtest.slide()
        after_value = self.stock_backtest.get_value()
        self.value_record.append(after_value)
        self.all_value_record.append(
            {"date": self.stock_backtest.get_cur_time(), "value": after_value}
        )
        if reward is None:
            if (delta := after_value - self.value_record[-2]) > 0:
                reward = max(delta / now_value * 100, 2)
            else:
                reward = min(20 * delta / now_value * 100, -20)

        return self.observation(), reward, self.done, False, {}

    def observation(self):
        all_value = self.stock_backtest.get_value()
        positions = self.stock_backtest.get_positions()

        features = [
            [
                self.stock_backtest.get_data(name, "open", 0, 0) / MAX_SHARE_PRICE,
                self.stock_backtest.get_data(name, "high", 0, 0) / MAX_SHARE_PRICE,
                self.stock_backtest.get_data(name, "low", 0, 0) / MAX_SHARE_PRICE,
                self.stock_backtest.get_data(name, "close", 0, 0) / MAX_SHARE_PRICE,
                self.stock_backtest.get_data(name, "volume", 0, 0) / MAX_VOLUME,
                self.stock_backtest.get_data(name, "amount", 0, 0) / MAX_AMOUNT,
                self.stock_backtest.get_data(name, "adjustflag", 0, 0) / 10,
                self.stock_backtest.get_data(name, "tradestatus", 0, 0) / 1,
                # self.stock_backtest.get_data(name, 'pctChg', 0) / 100,
                self.stock_backtest.get_data(name, "peTTM", 0, 0) / 1e4,
                self.stock_backtest.get_data(name, "pbMRQ", 0, 0) / 100,
                self.stock_backtest.get_data(name, "psTTM", 0, 0) / 100,
                self.stock_backtest.get_data(name, "pctChg", 0, 0) / 1e3,
                positions.get(name, PositionType(name, 0, 0)).size
                * self.stock_backtest.get_data(name, "close", 0, 0)
                / all_value,
                # positions.get(name, PositionType(name, 0, 0)).size / MAX_VOLUME
            ]
            for name in self.stock_backtest.datas
        ]
        result = np.array(features, dtype=np.float32)
        result = np.nan_to_num(result, nan=0)
        return result

    def render(self):
        print(
            f"Current state: {self.stock_backtest.get_cur_time()},\t {self.stock_backtest.get_value():.2f},\t "
            f"{self.stock_backtest.get_stock_value():.2f},\t {self.stock_backtest.get_cash():.2f}, \t"
            f"{self.stock_backtest.get_positions()}"
        )  # 打印当前状态

    def save_result(self, save_dir: str):
        df_result = pd.DataFrame(self.all_value_record)
        df_result.rename(columns={"date": "datetime", "value": "cash"}, inplace=True)
        result_file = os.path.join(save_dir, "cash_record.csv")
        df_result.to_csv(result_file)
        calculate_backtrade_result(result_file)


def get_backtest_model(start_date, end_date, stock_list):
    import baostock as bs

    backtest = BacktestModel()
    bs.login()
    # stock_list = get_baostock_list("2024-09-03").iloc[300:400]["code"]
    # stock_list = ["sh.600036"]
    for code in stock_list:
        df = get_baostock_data(code, start_date, end_date, is_cover=False)
        cols = [
            "open",
            "high",
            "low",
            "close",
            "volume",
            "amount",
            "adjustflag",
            "turn",
            "tradestatus",
            "pctChg",
            "peTTM",
            "pbMRQ",
            "psTTM",
            "pcfNcfTTM",
            "isST",
        ]
        df[cols] = df[cols].replace("", "0").astype(float)
        df.set_index("date", inplace=True)
        if len(df) > 0:
            backtest.add_data(df, code)
    bs.logout()
    backtest.set_backtest_date(start_date, end_date)
    backtest.set_commission(0, 0)
    backtest.set_origin_cash(10000000)
    return backtest


def get_shang_100():
    from xtquant.xtdata import get_stock_list_in_sector

    stocks = get_stock_list_in_sector("沪深300")
    new_stocks = []
    for x in stocks:
        mm = x.split(".")
        if mm[1] == "SH":
            new_stocks.append(f"sh.{mm[0]}")
        elif mm[1] == "SZ":
            new_stocks.append(f"sz.{mm[0]}")
    return new_stocks


def job_rl(
    train_start_date: str,
    train_end_date: str,
    test_start_date: str,
    test_end_date: str,
    codes: list,
    rl_model: str,
    policy: str,
    base_dir: str,
):
    """
    :param start_date: like 2016-01-01
    :param end_date: like 2016-01-01
    :param codes: ["sh.600001",...]
    :param save_path:
    :return:
    """
    save_path = os.path.join(base_dir, rl_model, policy)
    os.makedirs(save_path, exist_ok=True)
    info_file = os.path.join(save_path, "info.json")
    with open(info_file, "w") as f:
        json.dump(
            {
                "train_start_date": train_start_date,
                "train_end_date": train_end_date,
                "test_start_date": test_start_date,
                "test_end_date": test_end_date,
                "rl_model": rl_model,
                "policy": policy,
                "codes": codes,
            },
            f,
        )
    model_map = dict(PPO=PPO, SAC=SAC, A2C=A2C, DQN=DQN, TD3=TD3, DDPG=DDPG)
    model_type = model_map[rl_model]
    model_file = os.path.join(save_path, "model")
    backtest = get_backtest_model(train_start_date, train_end_date, codes)
    env = StockEnv(backtest)
    model = model_type(policy, env, verbose=1)
    model.learn(total_timesteps=20000)
    model.save(model_file)
    backtest = get_backtest_model(test_start_date, test_end_date, codes)
    env = StockEnv(backtest)
    model = model_type(policy, env, verbose=1)
    model.load(model_file)
    obs, info = env.reset()
    for i in range(10000):
        env.render()
        action, _states = model.predict(obs)
        obs, reward, done, _, info = env.step(action)
        if done:
            break
    env.save_result(save_path)
    env.close()


if __name__ == "__main__":
    pass

    # rl_model = sys.argv[1]
    rl_model = "PPO"
    codes = get_shang_100()
    # policy = "CnnPolicy"
    policy = "MlpPolicy"
    job_rl(
        "2020-01-01",
        "2023-06-30",
        "2023-07-01",
        "2024-09-30",
        codes,
        rl_model,
        policy,
        "rl_data_new10",
    )

    # job_rl(
    #     "2019-01-01",
    #     "2023-06-01",
    #     "2023-07-01",
    #     "2024-09-30",
    #     codes,
    #     rl_model,
    #     policy,
    #     "rl_data3"
    # )
